package cmd

import (
	//"context"
	"fmt"
	//"io/ioutil"
	//"strconv"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/spf13/cobra"

	"github.com/pingcap/tidb/config"
	"github.com/pingcap/tidb/store/tikv"
)

var restoreCmd = &cobra.Command{
	Use:   "restore",
	Short: "Restore He3DB",
	Long:  "Welcome to use hbr for He3DB restore",
	Run:   runRestore,
}

func init() {
	rootCmd.AddCommand(restoreCmd)
}

func runRestore(cmd *cobra.Command, args []string) {
	var sem = make(chan bool, concurrency)
	restoreStart := time.Now()
	access_key, _ := cmd.Flags().GetString("access_key")
	secret_key, _ := cmd.Flags().GetString("secret_key")
	endpoint, _ := cmd.Flags().GetString("endpoint")
	region, _ := cmd.Flags().GetString("region")
	bucket, _ := cmd.Flags().GetString("bucket")
	pd, _ := cmd.Flags().GetString("pd")
	backup_name, _ := cmd.Flags().GetString("name")

	if access_key == "" || secret_key == "" || endpoint == "" || region == "" || bucket == "" || pd == "" || backup_name == "" {
		fmt.Printf("PARAMETER ERROR!\n")
		return
	}

	client, err := tikv.NewRawKVClient([]string{pd}, config.Security{})
	if err != nil {
		fmt.Printf("Connect Tikv Error!\n%v\n", err)
		return
	}

	sess, err := session.NewSession(&aws.Config{
		Region:      aws.String(region),
		Endpoint:    aws.String(endpoint),
		Credentials: credentials.NewStaticCredentials(access_key, secret_key, ""),
		S3ForcePathStyle: aws.Bool(true),
	})
	if err != nil {
		fmt.Printf("Connect S3 Error!\n%v\n", err)
		return
	}
	s3_client := s3.New(sess)

	count := 0
	input := &s3.ListObjectsInput{
		Bucket: aws.String(bucket),
		Prefix: aws.String(backup_name),
	}
	for {
		resp, err := s3_client.ListObjects(input)
		if err != nil {
			fmt.Printf("S3 ListObjects Error!\n%v\n", err)
			return
		}

		for _, keys := range resp.Contents {
			wg.Add(1)
			sem <- true
			go s3RestoreKVRaw(s3_client, bucket, backup_name, keys, client, sem)
			count++
		}

		if resp.NextMarker == nil {
			fmt.Printf("Done!\n")
			break
		}
		input.Marker = resp.NextMarker
	}
	wg.Wait()
	fmt.Printf("N:%v\n", count)
	fmt.Println("restore time:", time.Since(restoreStart))
}


